import random
import time
from chess_pieces import Ju, Ma, Xiang, Shi, King, Pao, Pawn

class ChessAI:
    """中国象棋AI类，使用极小极大算法和Alpha-Beta剪枝"""
    
    def __init__(self, difficulty="medium", ai_color="black"):
        """初始化AI
        
        Args:
            difficulty (str): 难度级别 'easy', 'medium', 'hard'
            ai_color (str): AI执子颜色 'red' 或 'black'
        """
        self.difficulty = difficulty
        self.ai_color = ai_color
        self.transposition_table = {}  # 置换表
        
        # 根据难度级别设置参数
        self._set_difficulty_params()
        
        # 棋子基础价值
        self.piece_values = {
            "帅": 10000, "将": 10000,
            "車": 900, "车": 900,
            "馬": 400, "马": 400,
            "相": 200, "象": 200,
            "仕": 200, "士": 200,
            "炮": 450,
            "兵": 100, "卒": 100
        }
        
        # 位置价值表
        self._init_position_tables()
    
    def _set_difficulty_params(self):
        """根据难度设置AI参数"""
        if self.difficulty == "easy":
            self.search_depth = 2  # 增加搜索深度
            self.randomness = 0.15  # 减少随机性
            self.max_think_time = 1500  # 最多思考1.5秒
        elif self.difficulty == "medium":
            self.search_depth = 3  # 增加搜索深度
            self.randomness = 0.08  # 减少随机性
            self.max_think_time = 2500  # 最多思考2.5秒
        else:  # hard
            self.search_depth = 4  # 增加搜索深度
            self.randomness = 0.03  # 减少随机性
            self.max_think_time = 4000  # 最多思考4秒
        
        # 启用迭代加深搜索
        self.use_iterative_deepening = True
    
    def _init_position_tables(self):
        """初始化棋子位置价值表"""
        # 兵/卒位置价值表 (行, 列)
        self.pawn_pos_red = [
            [0,  0,  0,  0,  0,  0,  0,  0,  0],  # 0行
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [2,  4,  6,  6,  6,  6,  6,  4,  2],  # 过河后增值
            [2,  4,  6,  6, 10,  6,  6,  4,  2],
            [4,  8, 12, 12, 16, 12, 12,  8,  4],
            [6, 10, 16, 16, 22, 16, 16, 10,  6],
            [10, 14, 20, 20, 28, 20, 20, 14, 10]   # 9行（底线）
        ]
        
        # 黑方卒子位置价值表 (镜像红方兵)
        self.pawn_pos_black = list(reversed(self.pawn_pos_red))
        
        # 车位置价值表
        self.rook_pos_red = [
            [14, 14, 12, 18, 16, 18, 12, 14, 14],  # 0行
            [16, 20, 18, 24, 26, 24, 18, 20, 16],
            [12, 12, 12, 18, 18, 18, 12, 12, 12],
            [12, 18, 16, 22, 22, 22, 16, 18, 12],
            [12, 14, 12, 18, 18, 18, 12, 14, 12],
            [12, 16, 14, 20, 20, 20, 14, 16, 12],
            [12, 12, 12, 12, 12, 12, 12, 12, 12],
            [12, 18, 16, 22, 22, 22, 16, 18, 12],
            [12, 12, 12, 18, 18, 18, 12, 12, 12],
            [16, 22, 20, 24, 26, 24, 20, 22, 16]   # 9行（底线）
        ]
        
        self.rook_pos_black = list(reversed(self.rook_pos_red))
        
        # 马位置价值表
        self.knight_pos_red = [
            [4,  8,  16, 12, 4,  12, 16, 8,  4],   # 0行
            [4,  10, 28, 16, 8,  16, 28, 10, 4],
            [12, 14, 16, 20, 18, 20, 16, 14, 12],
            [8,  24, 18, 24, 20, 24, 18, 24, 8],
            [6,  16, 14, 18, 16, 18, 14, 16, 6],
            [4,  12, 16, 14, 12, 14, 16, 12, 4],
            [2,  6,  8,  6,  10, 6,  8,  6,  2],
            [4,  2,  8,  8,  4,  8,  8,  2,  4],
            [0,  2,  4,  4,  -2, 4,  4,  2,  0],
            [0,  -4, 0,  0,  0,  0,  0,  -4, 0]    # 9行（底线）
        ]
        
        self.knight_pos_black = list(reversed(self.knight_pos_red))
        
        # 炮位置价值表
        self.cannon_pos_red = [
            [6,  4,  0, -10, -12, -10, 0,  4,  6],  # 0行
            [2,  2,  0,  -4, -14, -4,  0,  2,  2],
            [2,  2,  0,  -10, -8, -10, 0,  2,  2],
            [0,  0, -2,   4,  10,  4, -2,  0,  0],
            [0,  0,  0,   2,   8,  2,  0,  0,  0],
            [4,  0,  8,   6,  10,  6,  8,  0,  4],
            [0,  0,  0,   2,   4,  2,  0,  0,  0],
            [4,  2,  8,  12,  12, 12,  8,  2,  4],
            [6,  6,  6,  12,  18, 12,  6,  6,  6],
            [6,  6,  0,   2,   4,  2,  0,  6,  6]   # 9行（底线）
        ]
        
        self.cannon_pos_black = list(reversed(self.cannon_pos_red))
        
        # 相/象位置价值表
        self.bishop_pos_red = [
            [0,  0,  0,  0,  0,  0,  0,  0,  0],   # 0行
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0, 20,  0,  0,  0, 20,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0, 20,  0,  0,  0,  0,  0, 20,  0]    # 9行（底线）
        ]
        
        self.bishop_pos_black = list(reversed(self.bishop_pos_red))
        
        # 士/仕位置价值表
        self.advisor_pos_red = [
            [0,  0,  0,  0,  0,  0,  0,  0,  0],   # 0行
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0, 20,  0, 20,  0,  0,  0],
            [0,  0,  0,  0, 23,  0,  0,  0,  0],
            [0,  0,  0, 20,  0, 20,  0,  0,  0]    # 9行（底线）
        ]
        
        self.advisor_pos_black = list(reversed(self.advisor_pos_red))
        
        # 将/帅位置价值表
        self.king_pos_red = [
            [0,  0,  0,  0,  0,  0,  0,  0,  0],   # 0行
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  0,  0,  0,  0,  0,  0],
            [0,  0,  0,  1,  1,  1,  0,  0,  0],
            [0,  0,  0,  2,  2,  2,  0,  0,  0],
            [0,  0,  0, 11, 15, 11,  0,  0,  0]    # 9行（底线）
        ]
        
        self.king_pos_black = list(reversed(self.king_pos_red))
        
    def get_move(self, game_state):
        """获取AI的最佳走法
        
        Args:
            game_state: GameState对象，表示当前棋盘状态
            
        Returns:
            tuple: ((from_row, from_col), (to_row, to_col)) 表示移动的起点和终点
        """
        # 清空置换表
        self.transposition_table = {}
        
        # 添加一点随机性
        if random.random() < self.randomness:
            return self._get_random_move(game_state)
        
        # 获取所有可能的走法
        valid_moves = self._get_valid_moves(game_state, self.ai_color)
        
        if not valid_moves:
            return None  # 无有效走法
        
        # 对走法进行排序（启发式）以提高剪枝效率
        valid_moves = self._sort_moves(game_state, valid_moves)
        
        # 如果只有一个有效移动，直接返回
        if len(valid_moves) == 1:
            return valid_moves[0]
        
        # 记录开始时间
        start_time = time.time()
        
        best_move = valid_moves[0]  # 默认使用第一个有效走法
        best_value = float('-inf')
        
        # 使用迭代加深搜索
        if self.use_iterative_deepening:
            # 从较浅的深度开始搜索，逐步加深
            for current_depth in range(1, self.search_depth + 1):
                # 检查是否超时
                if (time.time() - start_time) * 1000 > self.max_think_time * 0.7:
                    break
                
                alpha = float('-inf')
                beta = float('inf')
                current_best_move = None
                current_best_value = float('-inf')
                
                # 搜索当前深度的最佳走法
                for from_pos, to_pos in valid_moves:
                    # 检查思考时间是否超出限制
                    if (time.time() - start_time) * 1000 > self.max_think_time:
                        break
                    
                    # 模拟移动
                    cloned_state = self._clone_game_state(game_state)
                    self._make_move(cloned_state, from_pos, to_pos)
                    
                    # 获取走法评分
                    value = self._alpha_beta(cloned_state, current_depth - 1, 
                                          alpha, beta, False, start_time)
                    
                    # 更新最佳走法
                    if value > current_best_value:
                        current_best_value = value
                        current_best_move = (from_pos, to_pos)
                        alpha = max(alpha, current_best_value)
                
                # 更新全局最佳走法
                if current_best_move and current_best_value > best_value:
                    best_value = current_best_value
                    best_move = current_best_move
        else:
            # 原始的固定深度搜索
            alpha = float('-inf')
            beta = float('inf')
            
            for from_pos, to_pos in valid_moves:
                # 检查思考时间是否超出限制
                if (time.time() - start_time) * 1000 > self.max_think_time:
                    break
                
                # 模拟移动
                cloned_state = self._clone_game_state(game_state)
                self._make_move(cloned_state, from_pos, to_pos)
                
                # 获取走法评分
                value = self._alpha_beta(cloned_state, self.search_depth - 1, 
                                      alpha, beta, False, start_time)
                
                # 更新最佳走法
                if value > best_value:
                    best_value = value
                    best_move = (from_pos, to_pos)
                    alpha = max(alpha, best_value)
        
        return best_move
    
    def _get_random_move(self, game_state):
        """获取一个随机有效走法"""
        valid_moves = self._get_valid_moves(game_state, self.ai_color)
        if valid_moves:
            return random.choice(valid_moves)
        return None
    
    def _get_valid_moves(self, game_state, color):
        """获取指定颜色棋子的所有有效走法"""
        valid_moves = []
        
        for piece in game_state.pieces:
            if piece.color == color:
                # 获取该棋子所有可能的移动位置
                possible_moves, _ = game_state.calculate_possible_moves(piece.row, piece.col)
                
                # 添加到有效走法列表
                for to_row, to_col in possible_moves:
                    valid_moves.append(((piece.row, piece.col), (to_row, to_col)))
        
        return valid_moves
    
    def _sort_moves(self, game_state, moves):
        """启发式走法排序，提高剪枝效率
        
        排序优先级：吃子 > 将军 > 普通走法
        """
        scored_moves = []
        
        for from_pos, to_pos in moves:
            score = 0
            from_row, from_col = from_pos
            to_row, to_col = to_pos
            
            # 检查是否吃子
            target_piece = None
            for piece in game_state.pieces:
                if piece.row == to_row and piece.col == to_col:
                    target_piece = piece
                    break
            
            if target_piece:
                # 吃子得分：目标棋子价值 - 当前棋子价值/10
                moving_piece = None
                for piece in game_state.pieces:
                    if piece.row == from_row and piece.col == from_col:
                        moving_piece = piece
                        break
                
                score += self._get_piece_value(target_piece)
                score -= self._get_piece_value(moving_piece) // 10
            
            # 模拟移动，检查是否将军
            cloned_state = self._clone_game_state(game_state)
            self._make_move(cloned_state, from_pos, to_pos)
            
            opponent_color = "red" if self.ai_color == "black" else "black"
            if self._is_check(cloned_state, opponent_color):
                score += 200  # 将军得高分
            
            scored_moves.append(((from_pos, to_pos), score))
        
        # 按分数降序排列
        scored_moves.sort(key=lambda x: x[1], reverse=True)
        
        # 返回排序后的走法
        return [move for move, _ in scored_moves]
    
    def _alpha_beta(self, game_state, depth, alpha, beta, is_maximizing, start_time):
        """Alpha-Beta剪枝搜索
        
        Args:
            game_state: 游戏状态
            depth: 当前搜索深度
            alpha: Alpha值
            beta: Beta值
            is_maximizing: 是否是最大化层(AI回合)
            start_time: 搜索开始时间
            
        Returns:
            int: 局面评分
        """
        # 生成棋盘状态的唯一键
        state_key = self._get_state_key(game_state)
        
        # 检查置换表
        if state_key in self.transposition_table and self.transposition_table[state_key]['depth'] >= depth:
            return self.transposition_table[state_key]['value']
        
        # 达到叶节点或游戏结束
        if depth == 0 or game_state.game_over:
            value = self._evaluate_board(game_state)
            self.transposition_table[state_key] = {'value': value, 'depth': depth}
            return value
        
        if is_maximizing:  # AI回合(最大化)
            max_eval = float('-inf')
            moves = self._get_valid_moves(game_state, self.ai_color)
            moves = self._sort_moves(game_state, moves)
            
            for from_pos, to_pos in moves:
                # 检查思考时间是否超出限制
                if (time.time() - start_time) * 1000 > self.max_think_time:
                    # print("AI思考时间超限，提前返回")  # 可用于调试
                    break
                
                cloned_state = self._clone_game_state(game_state)
                self._make_move(cloned_state, from_pos, to_pos)
                
                eval = self._alpha_beta(cloned_state, depth - 1, alpha, beta, False, start_time)
                max_eval = max(max_eval, eval)
                alpha = max(alpha, eval)
                
                if beta <= alpha:
                    break
            
            self.transposition_table[state_key] = {'value': max_eval, 'depth': depth}
            return max_eval
        else:  # 玩家回合(最小化)
            min_eval = float('inf')
            opponent_color = "red" if self.ai_color == "black" else "black"
            moves = self._get_valid_moves(game_state, opponent_color)
            moves = self._sort_moves(game_state, moves)
            
            for from_pos, to_pos in moves:
                # 检查思考时间是否超出限制
                if (time.time() - start_time) * 1000 > self.max_think_time:
                    # print("AI思考时间超限，提前返回")  # 可用于调试
                    break
                
                cloned_state = self._clone_game_state(game_state)
                self._make_move(cloned_state, from_pos, to_pos)
                
                eval = self._alpha_beta(cloned_state, depth - 1, alpha, beta, True, start_time)
                min_eval = min(min_eval, eval)
                beta = min(beta, eval)
                
                if beta <= alpha:
                    break
            
            self.transposition_table[state_key] = {'value': min_eval, 'depth': depth}
            return min_eval
    
    def _evaluate_board(self, game_state):
        """评估局面分数
        
        返回正分对AI有利，负分对玩家有利
        """
        # 如果游戏已结束
        if game_state.game_over:
            if game_state.winner == self.ai_color:
                return 10000  # AI获胜
            else:
                return -10000  # 玩家获胜
        
        ai_score = 0
        player_score = 0
        
        # 计算棋子基础分值和位置加成
        for piece in game_state.pieces:
            piece_value = self._get_piece_value(piece)
            position_value = self._get_position_value(piece)
            
            # 根据棋子威胁情况调整价值
            threat_value = self._evaluate_piece_threats(game_state, piece)
            protection_value = self._evaluate_piece_protection(game_state, piece)
            
            # 计算总价值
            total_value = piece_value + position_value + protection_value - threat_value
            
            if piece.color == self.ai_color:
                ai_score += total_value
            else:
                player_score += total_value
        
        # 计算行动力和控制力
        ai_moves = self._get_valid_moves(game_state, self.ai_color)
        player_moves = self._get_valid_moves(game_state, "red" if self.ai_color == "black" else "black")
        
        # 行动力评估
        ai_mobility = len(ai_moves) * 3  # 增加行动力权重
        player_mobility = len(player_moves) * 3
        
        # 中心控制评估
        ai_center_control = self._evaluate_center_control(game_state, self.ai_color)
        player_center_control = self._evaluate_center_control(game_state, "red" if self.ai_color == "black" else "black")
        
        ai_score += ai_mobility + ai_center_control
        player_score += player_mobility + player_center_control
        
        # 检查将军状态
        if self._is_check(game_state, self.ai_color):
            ai_score -= 80  # AI被将军，扣分更多
        
        if self._is_check(game_state, "red" if self.ai_color == "black" else "black"):
            player_score -= 80  # 玩家被将军，扣分更多
            
        # 评估王的安全
        ai_king_safety = self._evaluate_king_safety(game_state, self.ai_color)
        player_king_safety = self._evaluate_king_safety(game_state, "red" if self.ai_color == "black" else "black")
        
        ai_score += ai_king_safety
        player_score += player_king_safety
        
        # 返回最终分差
        return ai_score - player_score
    
    def _get_piece_value(self, piece):
        """获取棋子基础价值"""
        if not piece:
            return 0
        return self.piece_values.get(piece.name, 0)
    
    def _get_position_value(self, piece):
        """获取棋子在特定位置的附加价值"""
        if not piece:
            return 0
            
        row, col = piece.row, piece.col
        
        # 根据棋子类型选择相应的位置价值表
        if isinstance(piece, Pawn):
            if piece.color == "red":
                return self.pawn_pos_red[row][col]
            else:
                return self.pawn_pos_black[row][col]
        
        elif isinstance(piece, Ju):
            if piece.color == "red":
                return self.rook_pos_red[row][col]
            else:
                return self.rook_pos_black[row][col]
        
        elif isinstance(piece, Ma):
            if piece.color == "red":
                return self.knight_pos_red[row][col]
            else:
                return self.knight_pos_black[row][col]
        
        elif isinstance(piece, Pao):
            if piece.color == "red":
                return self.cannon_pos_red[row][col]
            else:
                return self.cannon_pos_black[row][col]
        
        elif isinstance(piece, Xiang):
            if piece.color == "red":
                return self.bishop_pos_red[row][col]
            else:
                return self.bishop_pos_black[row][col]
        
        elif isinstance(piece, Shi):
            if piece.color == "red":
                return self.advisor_pos_red[row][col]
            else:
                return self.advisor_pos_black[row][col]
        
        elif isinstance(piece, King):
            if piece.color == "red":
                return self.king_pos_red[row][col]
            else:
                return self.king_pos_black[row][col]
        
        return 0
    
    def _is_check(self, game_state, color):
        """检查指定颜色是否被将军"""
        # 这里复用游戏规则已有的检查
        return color == game_state.player_turn and game_state.is_check
    
    def _clone_game_state(self, game_state):
        """创建游戏状态的深拷贝用于模拟"""
        # 简单实现，复制必要的状态
        from copy import deepcopy
        import types
        
        # 创建一个新的game_state对象
        cloned_state = types.SimpleNamespace()
        
        # 使用浅拷贝加速 - 我们只修改pieces，其余共享
        cloned_state.pieces = deepcopy(game_state.pieces)
        cloned_state.player_turn = game_state.player_turn
        cloned_state.game_over = game_state.game_over
        cloned_state.winner = game_state.winner
        cloned_state.is_check = game_state.is_check
        
        # 添加必要的方法 - 使用闭包以避免重复创建lambda
        def get_piece_at(row, col):
            for p in cloned_state.pieces:
                if p.row == row and p.col == col:
                    return p
            return None
        
        cloned_state.get_piece_at = get_piece_at
        cloned_state.calculate_possible_moves = game_state.calculate_possible_moves
        
        return cloned_state
    
    def _make_move(self, game_state, from_pos, to_pos):
        """在克隆的游戏状态中执行移动"""
        from_row, from_col = from_pos
        to_row, to_col = to_pos
        
        # 找出要移动的棋子
        moving_piece = None
        for piece in game_state.pieces:
            if piece.row == from_row and piece.col == from_col:
                moving_piece = piece
                break
        
        if not moving_piece:
            return False
        
        # 查找目标位置是否有棋子（吃子）
        target_piece = None
        for piece in game_state.pieces:
            if piece.row == to_row and piece.col == to_col:
                target_piece = piece
                break
        
        # 如果有目标棋子，从列表中移除
        if target_piece:
            game_state.pieces.remove(target_piece)
        
        # 更新棋子位置
        moving_piece.row = to_row
        moving_piece.col = to_col
        
        # 切换回合
        game_state.player_turn = "red" if game_state.player_turn == "black" else "black"
        
        return True
    
    def _get_state_key(self, game_state):
        """生成棋盘状态的唯一键，用于置换表"""
        # 简单实现，使用棋子位置组合
        pieces_str = ""
        for piece in sorted(game_state.pieces, key=lambda p: (p.color, p.name, p.row, p.col)):
            pieces_str += f"{piece.name}:{piece.color}:{piece.row}:{piece.col}|"
        
        return pieces_str + game_state.player_turn 

    def _evaluate_piece_threats(self, game_state, piece):
        """评估棋子受到的威胁"""
        if not piece:
            return 0
        
        threat_value = 0
        opponent_color = "red" if piece.color == "black" else "black"
        
        # 检查对方每个棋子是否可以吃掉当前棋子
        for opp_piece in game_state.pieces:
            if opp_piece.color == opponent_color:
                # 检查是否可以直接吃掉
                if self._can_capture(game_state, opp_piece, piece):
                    # 威胁值与威胁棋子价值和被威胁棋子价值相关
                    threat_factor = 0.2  # 基础威胁因子
                    
                    # 如果威胁棋子价值低于被威胁棋子，威胁更严重
                    if self._get_piece_value(opp_piece) < self._get_piece_value(piece):
                        threat_factor = 0.5
                        
                    threat_value += self._get_piece_value(piece) * threat_factor
        
        return threat_value

    def _evaluate_piece_protection(self, game_state, piece):
        """评估棋子受到的保护"""
        if not piece:
            return 0
        
        protection_value = 0
        
        # 检查己方每个棋子是否可以保护当前棋子
        for friendly_piece in game_state.pieces:
            if friendly_piece != piece and friendly_piece.color == piece.color:
                # 检查是否可以保护（能够吃掉吃掉当前棋子的对方棋子）
                if self._can_protect(game_state, friendly_piece, piece):
                    protection_value += self._get_piece_value(piece) * 0.1
        
        return protection_value

    def _can_capture(self, game_state, attacker, target):
        """检查攻击棋子是否可以吃掉目标棋子"""
        # 使用游戏规则检查移动是否合法
        from game_rules import GameRules
        return GameRules.is_valid_move(game_state.pieces, attacker, attacker.row, attacker.col, target.row, target.col)

    def _can_protect(self, game_state, protector, protected_piece):
        """检查保护棋子是否可以保护被保护棋子（能吃掉吃掉被保护棋子的对方棋子）"""
        # 简化实现：如果保护棋子可以移动到被保护棋子位置，则认为可以保护
        from game_rules import GameRules
        
        # 创建一个假想的对方棋子在被保护棋子位置
        opponent_color = "red" if protector.color == "black" else "black"
        
        # 检查保护棋子是否可以移动到被保护棋子位置
        # 这是简化的保护检查，实际上应该更复杂
        return GameRules.is_valid_move(game_state.pieces, protector, protector.row, protector.col, protected_piece.row, protected_piece.col)

    def _evaluate_center_control(self, game_state, color):
        """评估对中心区域的控制"""
        # 定义中心区域 (行3-6, 列3-5)
        center_rows = [3, 4, 5, 6]
        center_cols = [3, 4, 5]
        
        control_score = 0
        
        # 计算在中心区域的棋子数量
        for piece in game_state.pieces:
            if piece.color == color:
                # 棋子在中心区域得分
                if piece.row in center_rows and piece.col in center_cols:
                    # 根据棋子类型给予不同分数
                    if isinstance(piece, Pawn):
                        control_score += 5
                    elif isinstance(piece, Pao):
                        control_score += 8
                    elif isinstance(piece, Ma):
                        control_score += 10
                    elif isinstance(piece, Ju):
                        control_score += 12
                    else:
                        control_score += 3
                
                # 计算棋子对中心的控制（可以移动到中心的棋子）
                moves, _ = game_state.calculate_possible_moves(piece.row, piece.col)
                for move_row, move_col in moves:
                    if move_row in center_rows and move_col in center_cols:
                        control_score += 2
        
        return control_score

    def _evaluate_king_safety(self, game_state, color):
        """评估王的安全性"""
        # 找出王
        king = None
        for piece in game_state.pieces:
            if isinstance(piece, King) and piece.color == color:
                king = piece
                break
        
        if not king:
            return -500  # 没有王，非常危险
        
        safety_score = 0
        
        # 检查周围是否有保护棋子
        protectors = 0
        for piece in game_state.pieces:
            if piece.color == color and not isinstance(piece, King):
                # 计算与王的距离
                distance = abs(piece.row - king.row) + abs(piece.col - king.col)
                if distance <= 2:  # 在王附近
                    protectors += 1
                    safety_score += 10
        
        # 王在九宫格中央更安全
        if color == "red":
            if king.row == 9 and king.col == 4:  # 红方王在九宫中央
                safety_score += 30
        else:  # black
            if king.row == 0 and king.col == 4:  # 黑方王在九宫中央
                safety_score += 30
        
        return safety_score 